import py_trees
import time


# 行为节点
class checkBattery(py_trees.behaviour.Behaviour):
    def __init__(self, name):
        super().__init__(name)
        self.battery_level = 50

    def update(self):
        if self.battery_level > 20:
            print(f"{self.name} 电池电量充足: {self.battery_level}")
            return py_trees.common.Status.SUCCESS
        else:
            print(f"{self.name} 电池电量不足: {self.battery_level}")
            return py_trees.common.Status.FAILURE

    def set_battery_level(self, level):
        self.battery_level = level


class MoveRobot(py_trees.behaviour.Behaviour):
    def __init__(self, name, target):
        super().__init__(name)
        self.name = name
        self.target = target

    def update(self):
        print(f"{self.name} 移动机器执行: {self.target}")
        return py_trees.common.Status.SUCCESS


# 创建行为树
def create_behavior_tree(battery_level):
    # 根节点
    root = py_trees.composites.Sequence("root根序列", memory=True)
    # 行为节点
    check_battery = checkBattery("battery")
    check_battery.set_battery_level(battery_level)

    move_robot = MoveRobot("机器人", "执行前进移动")
    # 构建行为树
    root.add_children(
        [
            check_battery,
            move_robot,
        ]
    )
    return root


def main():
    print("第1次运行行为树 ----------------------------")
    # 创建
    behaviour_tree = create_behavior_tree(50)
    # 打印树结构
    print("行为树结构:")
    print(py_trees.display.ascii_tree(behaviour_tree))

    # 运行
    behaviour_tree.tick_once()
    time.sleep(1)

    print("第2次运行行为树 ----------------------------")
    # 创建
    behaviour_tree = create_behavior_tree(10)
    # 打印树结构
    print("行为树结构:")
    print(py_trees.display.ascii_tree(behaviour_tree))

    # 运行
    behaviour_tree.tick_once()
    time.sleep(1)


if __name__ == "__main__":
    main()

"""运行结果
第1次运行行为树 ----------------------------
battery 电池电量充足: 50
机器人 移动机器执行: 执行前进移动

第2次运行行为树 ----------------------------
battery 电池电量不足: 10

"""
